package com.example.base;

import com.example.entity.Order;
import com.example.entity.Payment;
import org.apache.flink.api.common.functions.JoinFunction;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.ProcessJoinFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * Author wangJinLong
 * Date 2025/8/15 09:49
 **/
public class FlinkStreamJoin {
    public static void main(String[] args) throws Exception {
        slidingJoin();

    }

    private static void intervalJoin() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(3);

        DataStream<Order> mapSource = FlinkStreamSource.getOrderSource(environment);

        DataStream<Payment> paymentSource = FlinkStreamSource.getPaymentSource(environment);

        DataStream<Payment> process = mapSource.keyBy((KeySelector<Order, String>) Order::getOrderId)
                .intervalJoin(paymentSource.keyBy((KeySelector<Payment, String>) Payment::getOrderId))
                .between(Duration.ofSeconds(-5), Duration.ofSeconds(5))
                .process(new ProcessJoinFunction<>() {
                    @Override
                    public void processElement(Order left,
                                               Payment right,
                                               ProcessJoinFunction<Order, Payment, Payment>.Context ctx,
                                               Collector<Payment> out) throws Exception {
                        Payment payment = new Payment();
                        payment.setOrderId(left.getOrderId());
                        payment.setPayId(right.getPayId());
                        payment.setPayAmount(1);
                        out.collect(payment);
                    }
                });

        DataStream<Payment> payAmount =
                process.keyBy((KeySelector<Payment, String>) Payment::getPayId).sum("payAmount");

        payAmount.print();

        environment.execute();
    }

    private static void slidingJoin() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(3);

        DataStream<Order> mapSource = FlinkStreamSource.getOrderSource(environment);

        DataStream<Payment> paymentSource = FlinkStreamSource.getPaymentSource(environment);

        DataStream<Payment> dataStream = mapSource.join(paymentSource)
                .where(new KeySelector<Order, Object>() {
                    @Override
                    public Object getKey(Order value) throws Exception {
                        return value.getOrderId();
                    }
                })
                .equalTo(new KeySelector<Payment, Object>() {
                    @Override
                    public Object getKey(Payment value) throws Exception {
                        return value.getOrderId();
                    }
                })
                .window(SlidingEventTimeWindows.of(Duration.ofSeconds(10), Duration.ofSeconds(5)))
                .apply(new JoinFunction<Order, Payment, Payment>() {
                    @Override
                    public Payment join(Order first, Payment second) throws Exception {
                        Payment payment = new Payment();
                        payment.setOrderId(second.getOrderId());
                        payment.setPayId(second.getPayId());
                        payment.setPayAmount(first.getAmount());
                        return payment;
                    }
                });

        SingleOutputStreamOperator<Payment> payAmount = dataStream.keyBy(new KeySelector<Payment, Object>() {
            @Override
            public Object getKey(Payment value) throws Exception {
                return value.getPayId();
            }
        }).sum("payAmount");

        payAmount.print();

        environment.execute();
    }

    private static void tumbleJoin() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(3);

        DataStream<Order> mapSource = FlinkStreamSource.getOrderSource(environment);

        DataStream<Payment> paymentSource = FlinkStreamSource.getPaymentSource(environment);

        DataStream<Tuple2<String, Integer>> dataStream = mapSource.join(paymentSource)
                .where(new KeySelector<Order, Object>() {
                    @Override
                    public Object getKey(Order value) {
                        return value.getOrderId();
                    }
                })
                .equalTo(new KeySelector<Payment, Object>() {
                    @Override
                    public Object getKey(Payment value) {
                        return value.getOrderId();
                    }
                })
                .window(TumblingEventTimeWindows.of(Duration.ofSeconds(10)))
                .apply(new JoinFunction<Order, Payment, Tuple2<String, Integer>>() {
                    @Override
                    public Tuple2<String, Integer> join(Order order, Payment payment) throws Exception {
                        return new Tuple2<>(order.getOrderId(), 1);
                    }
                });

//        dataStream.print();

        DataStream<Tuple2<String, Integer>> operator =
                dataStream.keyBy((KeySelector<Tuple2<String, Integer>, Object>) value -> value.f0).sum(1);

        operator.print();

        environment.execute();
    }
}
